"""
RGB主题色控制器 - 优化版
改进：
1. 原生Windows托盘图标（高DPI支持）
2. 颜色简化和校正
3. 监听注册表变化（替代轮询）
4. Per-Monitor DPI V2 支持（清晰菜单）
"""

import colorsys
import ctypes
import os
import sys
import threading
import time
import winreg
from ctypes import wintypes

import win32api
import win32con
import win32gui
from openrgb import OpenRGBClient
from openrgb.utils import RGBColor

# 设置DPI感知 - 必须在创建任何窗口之前调用
try:
    # Windows 10 1703+ (Per-Monitor DPI V2 - 最佳)
    ctypes.windll.shcore.SetProcessDpiAwareness(2)  # PROCESS_PER_MONITOR_DPI_AWARE_V2
except:
    try:
        # Windows 8.1+ (Per-Monitor DPI V1)
        ctypes.windll.shcore.SetProcessDpiAwareness(1)
    except:
        try:
            # Windows Vista+ (System DPI Aware - 降级方案)
            ctypes.windll.user32.SetProcessDPIAware()
        except:
            pass  # 如果都失败，使用默认DPI设置

# Windows常量定义
WM_USER = 0x0400
WM_TRAYICON = WM_USER + 1
NIM_ADD = 0
NIM_MODIFY = 1
NIM_DELETE = 2
NIF_ICON = 2
NIF_MESSAGE = 1
NIF_TIP = 4
WM_LBUTTONDBLCLK = 0x0203
WM_RBUTTONUP = 0x0205

class ColorQuantizer:
    """颜色量化器 - 简化颜色到标准色"""

    # 标准色彩表（更好的RGB显示效果）
    STANDARD_COLORS = {
        'red': (255, 0, 0),
        'orange': (255, 127, 0),
        'yellow': (255, 255, 0),
        'lime': (127, 255, 0),
        'green': (0, 255, 0),
        'cyan': (0, 255, 255),
        'blue': (0, 0, 255),
        'purple': (127, 0, 255),
        'magenta': (255, 0, 255),
        'pink': (255, 0, 127),
        'white': (255, 255, 255),
    }

    @staticmethod
    def color_distance(c1, c2):
        """计算两个颜色的感知距离"""
        # 使用加权欧氏距离（更符合人眼感知）
        r_mean = (c1[0] + c2[0]) / 2
        r_diff = c1[0] - c2[0]
        g_diff = c1[1] - c2[1]
        b_diff = c1[2] - c2[2]

        return ((2 + r_mean/256) * r_diff**2 +
                4 * g_diff**2 +
                (2 + (255-r_mean)/256) * b_diff**2) ** 0.5

    @staticmethod
    def quantize_to_nearest(r, g, b, threshold=80):
        """量化到最近的标准颜色"""
        color = (r, g, b)

        # 找到最接近的标准色
        min_distance = float('inf')
        nearest_color = color
        nearest_name = 'custom'

        for name, std_color in ColorQuantizer.STANDARD_COLORS.items():
            distance = ColorQuantizer.color_distance(color, std_color)
            if distance < min_distance:
                min_distance = distance
                nearest_color = std_color
                nearest_name = name

        # 如果距离小于阈值，使用标准色，否则保持原色
        if min_distance < threshold:
            print(f"🎨 颜色简化: RGB({r},{g},{b}) → {nearest_name.upper()} RGB{nearest_color}")
            return nearest_color
        else:
            return color

    @staticmethod
    def adjust_saturation(r, g, b, saturation_boost=1.2):
        """增强饱和度（让RGB更鲜艳）"""
        # RGB转HSV
        h, s, v = colorsys.rgb_to_hsv(r/255, g/255, b/255)

        # 提升饱和度
        s = min(1.0, s * saturation_boost)

        # HSV转RGB
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        return int(r * 255), int(g * 255), int(b * 255)

    @staticmethod
    def process_color(r, g, b, quantize=True, boost_saturation=True):
        """处理颜色：量化+饱和度增强"""
        # 先增强饱和度
        if boost_saturation:
            r, g, b = ColorQuantizer.adjust_saturation(r, g, b)

        # 再量化到标准色
        if quantize:
            r, g, b = ColorQuantizer.quantize_to_nearest(r, g, b)

        return r, g, b

class RGBThemeController:
    def __init__(self):
        self.client = None
        self.devices = []
        self.current_color = None
        self.running = True
        self.display_on = True

        # 托盘相关
        self.hwnd = None
        self.menu = None

        # 颜色处理设置
        self.enable_quantize = True
        self.enable_saturation_boost = True

        # 注册表监听
        self.registry_thread = None

        print("🚀 正在启动RGB主题色控制器...")

    def connect_openrgb(self, max_retries=5):
        """连接到OpenRGB SDK服务器"""
        for attempt in range(max_retries):
            try:
                print(f"🔌 尝试连接OpenRGB (尝试 {attempt + 1}/{max_retries})...")
                self.client = OpenRGBClient(address='127.0.0.1', port=6742)
                self.devices = self.client.devices
                print(f"✅ 成功连接! 检测到 {len(self.devices)} 个设备:")
                for device in self.devices:
                    print(f"  - {device.name} ({len(device.leds)} LEDs)")
                return True
            except Exception as e:
                print(f"❌ 连接失败: {e}")
                if attempt < max_retries - 1:
                    print("⏳ 5秒后重试...")
                    time.sleep(5)

        print("⚠️ 无法连接到OpenRGB")
        return False

    def get_windows_accent_color(self):
        """从Windows注册表读取主题强调色"""
        try:
            key = winreg.OpenKey(
                winreg.HKEY_CURRENT_USER,
                r"Software\Microsoft\Windows\DWM"
            )

            accent_color, _ = winreg.QueryValueEx(key, "AccentColor")
            winreg.CloseKey(key)

            # 解析AABBGGRR格式
            r = accent_color & 0xFF
            g = (accent_color >> 8) & 0xFF
            b = (accent_color >> 16) & 0xFF

            return r, g, b

        except Exception as e:
            print(f"❌ 读取主题色失败: {e}")
            return 0, 120, 215

    def apply_color(self, r, g, b, process=True):
        """应用颜色到所有RGB设备"""
        if not self.client or not self.devices:
            return

        try:
            # 颜色处理
            if process:
                r, g, b = ColorQuantizer.process_color(
                    r, g, b,
                    quantize=self.enable_quantize,
                    boost_saturation=self.enable_saturation_boost
                )

            color = RGBColor(r, g, b)
            for device in self.devices:
                if "Sony DualSense" in device.name:
                    device.zones[0].set_color(color)
                    continue
                device.set_color(color)
            self.current_color = (r, g, b)
            print(f"✅ 应用颜色: RGB({r}, {g}, {b})")
        except Exception as e:
            print(f"❌ 应用颜色失败: {e}")
            self.reconnect_openrgb()

    def set_device_mode(self):
        """设置所有RGB设备的模式为Direct（如果支持）"""
        if not self.client or not self.devices:
            return

        try:
            for device in self.devices:
                device.set_mode('Direct')
        except Exception as e:
            print(f"❌ 设置模式失败: {e}")
            self.reconnect_openrgb()

    def reconnect_openrgb(self):
        """重新连接OpenRGB"""
        print("🔄 尝试重新连接OpenRGB...")
        if self.connect_openrgb(max_retries=3):
            if self.display_on and self.current_color:
                self.apply_color(*self.current_color, process=False)

    def turn_off_rgb(self):
        """关闭所有RGB（全黑）"""
        print("⚫ 显示器关闭，关闭RGB...")
        self.apply_color(0, 0, 0, process=False)
        self.display_on = False

    def restore_theme_color(self):
        """恢复Windows主题色"""
        print("🟢 显示器唤醒，恢复主题色...")
        r, g, b = self.get_windows_accent_color()
        self.apply_color(r, g, b)
        self.display_on = True

    def monitor_registry_changes(self):
        """监听注册表变化（替代轮询）"""
        print("👁️ 开始监听Windows主题色变化...")

        try:
            # 使用win32api打开注册表键
            import win32event

            key_handle = win32api.RegOpenKeyEx(
                win32con.HKEY_CURRENT_USER,
                r"Software\Microsoft\Windows\DWM",
                0,
                win32con.KEY_NOTIFY | win32con.KEY_READ
            )

            # 创建事件对象
            event = win32event.CreateEvent(None, False, False, None)

            while self.running:
                try:
                    # 注册通知（异步）
                    win32api.RegNotifyChangeKeyValue(
                        key_handle,
                        False,  # 不监听子键
                        win32api.REG_NOTIFY_CHANGE_LAST_SET,  # 监听值变化
                        event,
                        True    # 异步等待
                    )

                    # 等待事件触发（最多30秒超时，避免卡死）
                    result = win32event.WaitForSingleObject(event, 30000)

                    if result == win32event.WAIT_OBJECT_0 and self.running and self.display_on:
                        # 检测到变化，读取新颜色
                        time.sleep(0.3)  # 短暂延迟确保值已更新
                        r, g, b = self.get_windows_accent_color()

                        if self.current_color != (r, g, b):
                            print(f"🎨 检测到主题色变化!")
                            self.apply_color(r, g, b)

                except Exception as e:
                    if self.running:
                        print(f"⚠️ 注册表监听错误: {e}")
                        time.sleep(5)

            win32api.RegCloseKey(key_handle)
            win32event.CloseHandle(event)

        except Exception as e:
            print(f"❌ 无法启动注册表监听: {e}")
            print("⚠️ 回退到轮询模式...")
            # 回退到轮询
            self.monitor_theme_color_poll()

    def monitor_theme_color_poll(self):
        """轮询监控主题色（备用方案）"""
        while self.running:
            if self.display_on:
                r, g, b = self.get_windows_accent_color()
                if self.current_color != (r, g, b):
                    print(f"🎨 检测到主题色变化（轮询）")
                    self.apply_color(r, g, b)
            time.sleep(30)

    def create_native_icon(self):
        """创建原生高DPI托盘图标"""
        # 创建设备上下文
        hdc = win32gui.GetDC(0)
        hdc_mem = win32gui.CreateCompatibleDC(hdc)

        # 创建32x32位图（支持高DPI）
        size = 32
        bitmap = win32gui.CreateCompatibleBitmap(hdc, size, size)
        win32gui.SelectObject(hdc_mem, bitmap)

        # 背景透明（白色）
        brush_bg = win32gui.CreateSolidBrush(win32api.RGB(255, 255, 255))
        win32gui.SelectObject(hdc_mem, brush_bg)
        win32gui.Rectangle(hdc_mem, 0, 0, size, size)

        # 绘制渐变RGB圆环
        # 外圆 - 红色
        brush_outer = win32gui.CreateSolidBrush(win32api.RGB(255, 0, 0))
        pen_outer = win32gui.CreatePen(win32con.PS_SOLID, 3, win32api.RGB(255, 0, 0))
        win32gui.SelectObject(hdc_mem, brush_outer)
        win32gui.SelectObject(hdc_mem, pen_outer)
        win32gui.Ellipse(hdc_mem, 4, 4, 28, 28)

        # 中圆 - 绿色
        brush_mid = win32gui.CreateSolidBrush(win32api.RGB(0, 255, 0))
        pen_mid = win32gui.CreatePen(win32con.PS_SOLID, 2, win32api.RGB(0, 255, 0))
        win32gui.SelectObject(hdc_mem, brush_mid)
        win32gui.SelectObject(hdc_mem, pen_mid)
        win32gui.Ellipse(hdc_mem, 8, 8, 24, 24)

        # 内圆 - 蓝色
        brush_inner = win32gui.CreateSolidBrush(win32api.RGB(0, 100, 255))
        win32gui.SelectObject(hdc_mem, brush_inner)
        win32gui.Ellipse(hdc_mem, 12, 12, 20, 20)

        # 转换为图标
        icon_info = win32gui.CreateIconIndirect((
            True,  # fIcon
            0, 0,  # xHotspot, yHotspot
            bitmap,  # hbmMask
            bitmap   # hbmColor
        ))

        # 清理
        win32gui.DeleteObject(brush_bg)
        win32gui.DeleteObject(brush_outer)
        win32gui.DeleteObject(brush_mid)
        win32gui.DeleteObject(brush_inner)
        win32gui.DeleteObject(pen_outer)
        win32gui.DeleteObject(pen_mid)
        win32gui.DeleteObject(bitmap)
        win32gui.DeleteDC(hdc_mem)
        win32gui.ReleaseDC(0, hdc)

        return icon_info

    def create_window(self):
        """创建隐藏窗口（用于托盘和消息）"""
        wc = win32gui.WNDCLASS()
        wc.lpfnWndProc = self.wnd_proc
        wc.lpszClassName = "RGBControllerWindow"
        wc.hInstance = win32api.GetModuleHandle(None)
        wc.hCursor = win32gui.LoadCursor(0, win32con.IDC_ARROW)

        class_atom = win32gui.RegisterClass(wc)
        self.hwnd = win32gui.CreateWindow(
            class_atom,
            "RGB Controller",
            0, 0, 0, 0, 0, 0, 0,
            wc.hInstance, None
        )

        # 注册显示器电源通知
        GUID_CONSOLE_DISPLAY_STATE = "{6FE69556-704A-47A0-8F24-C28D936FDA47}"

        class GUID(ctypes.Structure):
            _fields_ = [
                ("Data1", wintypes.DWORD),
                ("Data2", wintypes.WORD),
                ("Data3", wintypes.WORD),
                ("Data4", wintypes.BYTE * 8)
            ]

        guid = GUID()
        ctypes.windll.ole32.CLSIDFromString(GUID_CONSOLE_DISPLAY_STATE, ctypes.byref(guid))
        ctypes.windll.user32.RegisterPowerSettingNotification(
            self.hwnd, ctypes.byref(guid), 0
        )

        # 添加托盘图标
        self.add_tray_icon()

        print("✅ 窗口和托盘图标创建成功")

    def add_tray_icon(self):
        """添加托盘图标"""
        icon = self.create_native_icon()

        nid = (
            self.hwnd,
            0,
            NIF_ICON | NIF_MESSAGE | NIF_TIP,
            WM_TRAYICON,
            icon,
            "RGB主题色控制器\n同步Windows主题色"
        )

        win32gui.Shell_NotifyIcon(NIM_ADD, nid)

    def show_menu(self):
        """显示右键菜单"""
        self.menu = win32gui.CreatePopupMenu()

        # 添加菜单项
        win32gui.AppendMenu(self.menu, win32con.MF_STRING, 1001,
                            f"当前颜色: RGB{self.current_color if self.current_color else '(无)'}")
        win32gui.AppendMenu(self.menu, win32con.MF_SEPARATOR, 0, "")

        # 颜色处理选项
        quantize_flag = win32con.MF_CHECKED if self.enable_quantize else win32con.MF_UNCHECKED
        win32gui.AppendMenu(self.menu, win32con.MF_STRING | quantize_flag, 1002,
                            "启用颜色简化")

        saturation_flag = win32con.MF_CHECKED if self.enable_saturation_boost else win32con.MF_UNCHECKED
        win32gui.AppendMenu(self.menu, win32con.MF_STRING | saturation_flag, 1003,
                            "增强饱和度")

        win32gui.AppendMenu(self.menu, win32con.MF_SEPARATOR, 0, "")

        # 开机启动选项
        startup_flag = win32con.MF_CHECKED if self.is_in_startup() else win32con.MF_UNCHECKED
        win32gui.AppendMenu(self.menu, win32con.MF_STRING | startup_flag, 1006,
                            "开机自启动")

        win32gui.AppendMenu(self.menu, win32con.MF_SEPARATOR, 0, "")
        win32gui.AppendMenu(self.menu, win32con.MF_STRING, 1004, "刷新颜色")
        win32gui.AppendMenu(self.menu, win32con.MF_STRING, 1005, "重新连接OpenRGB")
        win32gui.AppendMenu(self.menu, win32con.MF_SEPARATOR, 0, "")
        win32gui.AppendMenu(self.menu, win32con.MF_STRING, 1099, "退出")

        # 获取鼠标位置
        pos = win32gui.GetCursorPos()

        # 显示菜单（需要先设置为前台窗口）
        win32gui.SetForegroundWindow(self.hwnd)
        win32gui.TrackPopupMenu(
            self.menu, win32con.TPM_LEFTALIGN,
            pos[0], pos[1], 0, self.hwnd, None
        )
        win32gui.PostMessage(self.hwnd, win32con.WM_NULL, 0, 0)

    def is_in_startup(self):
        """检查是否在开机启动中"""
        try:
            key = winreg.OpenKey(
                winreg.HKEY_CURRENT_USER,
                r"Software\Microsoft\Windows\CurrentVersion\Run",
                0,
                winreg.KEY_READ
            )
            try:
                winreg.QueryValueEx(key, "RGBThemeController")
                winreg.CloseKey(key)
                return True
            except FileNotFoundError:
                winreg.CloseKey(key)
                return False
        except:
            return False

    def toggle_startup(self):
        """切换开机启动状态"""
        if self.is_in_startup():
            remove_from_startup()
        else:
            add_to_startup()

    def wnd_proc(self, hwnd, msg, wparam, lparam):
        """Windows消息处理"""
        # 托盘图标消息
        if msg == WM_TRAYICON:
            if lparam == WM_RBUTTONUP:
                self.show_menu()
            return 0

        # 菜单命令
        elif msg == win32con.WM_COMMAND:
            cmd_id = win32api.LOWORD(wparam)

            if cmd_id == 1002:  # 切换颜色简化
                self.enable_quantize = not self.enable_quantize
                status = "启用" if self.enable_quantize else "禁用"
                print(f"🎨 颜色简化: {status}")
                self.refresh_color()

            elif cmd_id == 1003:  # 切换饱和度增强
                self.enable_saturation_boost = not self.enable_saturation_boost
                status = "启用" if self.enable_saturation_boost else "禁用"
                print(f"🎨 饱和度增强: {status}")
                self.refresh_color()

            elif cmd_id == 1006:  # 切换开机启动
                self.toggle_startup()

            elif cmd_id == 1004:  # 刷新颜色
                self.refresh_color()

            elif cmd_id == 1005:  # 重新连接
                self.reconnect_openrgb()

            elif cmd_id == 1099:  # 退出
                self.quit_app()

            return 0

        # 电源管理消息
        elif msg == 0x0218:  # WM_POWERBROADCAST
            if wparam == 0x8013:  # PBT_POWERSETTINGCHANGE
                class POWERBROADCAST_SETTING(ctypes.Structure):
                    _fields_ = [
                        ("PowerSetting", ctypes.c_byte * 16),
                        ("DataLength", ctypes.c_uint32),
                        ("Data", ctypes.c_uint32)
                    ]

                setting = ctypes.cast(lparam, ctypes.POINTER(POWERBROADCAST_SETTING)).contents

                if setting.Data == 0:
                    self.turn_off_rgb()
                elif setting.Data == 1:
                    self.restore_theme_color()
            return 0

        elif msg == win32con.WM_DESTROY:
            win32gui.PostQuitMessage(0)
            return 0

        return win32gui.DefWindowProc(hwnd, msg, wparam, lparam)

    def refresh_color(self):
        """手动刷新颜色"""
        print("🔄 手动刷新颜色...")
        r, g, b = self.get_windows_accent_color()
        self.apply_color(r, g, b)

    def quit_app(self):
        """退出应用"""
        print("\n👋 正在退出...")
        self.running = False

        # 移除托盘图标
        nid = (self.hwnd, 0)
        win32gui.Shell_NotifyIcon(NIM_DELETE, nid)

        # 关闭RGB
        if self.client:
            self.apply_color(0, 0, 0, process=False)

        win32gui.DestroyWindow(self.hwnd)
        time.sleep(0.5)
        os._exit(0)

    def start(self):
        """启动控制器"""
        # 连接OpenRGB
        if not self.connect_openrgb():
            print("⚠️ OpenRGB连接失败，程序会继续运行并定期重试")

        self.set_device_mode()

        # 初始应用主题色
        r, g, b = self.get_windows_accent_color()
        print(f"🎨 当前Windows主题色: RGB({r}, {g}, {b})")
        self.apply_color(r, g, b)

        # 启动注册表监听线程
        self.registry_thread = threading.Thread(
            target=self.monitor_registry_changes,
            daemon=True
        )
        self.registry_thread.start()

        # 创建窗口和托盘
        self.create_window()

        print("\n✅ RGB主题色控制器运行中...")
        print("📊 功能状态:")
        print("  ✓ Windows主题色实时同步（注册表监听）")
        print("  ✓ 显示器开关检测")
        print("  ✓ 颜色简化和校正")
        print("  ✓ 原生高DPI托盘图标")
        print("\n右键托盘图标可进行操作\n")

        # 消息循环
        win32gui.PumpMessages()

def add_to_startup():
    """添加到开机启动"""
    try:
        key = winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Software\Microsoft\Windows\CurrentVersion\Run",
            0, winreg.KEY_SET_VALUE
        )

        # 检测是否为打包的exe
        if getattr(sys, 'frozen', False):
            # 打包后的exe
            exe_path = sys.executable
        else:
            # 开发模式（Python脚本）
            python_path = sys.executable.replace("python.exe", "pythonw.exe")
            script_path = os.path.abspath(__file__)
            exe_path = f'"{python_path}" "{script_path}"'

        winreg.SetValueEx(
            key, "RGBThemeController", 0, winreg.REG_SZ,
            exe_path if getattr(sys, 'frozen', False) else exe_path
        )

        winreg.CloseKey(key)
        print("✅ 已添加到开机启动")
        return True
    except Exception as e:
        print(f"❌ 添加开机启动失败: {e}")
        return False

def remove_from_startup():
    """从开机启动移除"""
    try:
        key = winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Software\Microsoft\Windows\CurrentVersion\Run",
            0, winreg.KEY_SET_VALUE
        )

        try:
            winreg.DeleteValue(key, "RGBThemeController")
            print("✅ 已从开机启动移除")
        except FileNotFoundError:
            print("ℹ️ 未在开机启动中找到")

        winreg.CloseKey(key)
        return True
    except Exception as e:
        print(f"❌ 移除失败: {e}")
        return False

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "--add-startup":
            add_to_startup()
            sys.exit(0)
        elif sys.argv[1] == "--remove-startup":
            remove_from_startup()
            sys.exit(0)

    controller = RGBThemeController()
    controller.start()